{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a9499b4e",
   "metadata": {},
   "source": [
    "## 记忆\n",
    "\n",
    "Langgraph 的记忆分为短期记忆和长期记忆。\n",
    "\n",
    "短期记忆是针对单个会话的，在会话中随时可以被调用。每次调用完图之后，会自动更新，然后在每个超步开始的时候读取。\n",
    "\n",
    "长期记忆是跨会话线程的，可以在任何时间与任何线程中调用。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "cfd51669",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "开始测试长期记忆系统\n",
      "\n",
      "Alice 第一次对话:\n",
      "记忆已保存: 我是Alice，我喜欢读书...\n",
      "使用 2 条消息作为上下文\n",
      "AI回复: 你好，Alice！很高兴认识你。你喜欢读书，那一定读过很多有趣的书吧？我特别想知道你最近在读什么书，或者有没有特别喜欢的类型？比如小说、科幻、历史还是其他类型的书呢？\n",
      "记忆已保存: 你好，Alice！很高兴认识你。你喜欢读书，那一定读过很多有...\n",
      "Alice 第二次对话:\n",
      "记忆已保存: 我刚才说喜欢什么？...\n",
      "使用 4 条消息作为上下文\n",
      "AI回复: 你刚才说：“我是Alice，我喜欢读书。” 所以你刚才说的是你喜欢读书。😊\n",
      "记忆已保存: 你刚才说：“我是Alice，我喜欢读书。” 所以你刚才说的是...\n",
      "\n",
      "Bob 第一次对话:\n",
      "记忆已保存: 我是Bob，我喜欢运动...\n",
      "使用 2 条消息作为上下文\n",
      "AI回复: 你好，Bob！很高兴认识你。喜欢运动是个非常好的习惯，能让人保持健康和活力。你平时都喜欢哪些运动呢？是喜欢跑步、打球，还是健身、游泳之类的？我很好奇你的运动经历和爱好！😊\n",
      "记忆已保存: 你好，Bob！很高兴认识你。喜欢运动是个非常好的习惯，能让人...\n",
      "Bob 第二次对话:\n",
      "记忆已保存: 我叫什么名字？...\n",
      "使用 4 条消息作为上下文\n",
      "AI回复: 你叫Bob呀！之前你已经告诉过我了，说\"我是Bob，我喜欢运动\"。所以，我当然知道你的名字是Bob啦！😊 你是不是在测试我的记忆力呢？不过没关系，我很乐意和你聊天，无论你问什么问题，我都会认真回答的！\n",
      "记忆已保存: 你叫Bob呀！之前你已经告诉过我了，说\"我是Bob，我喜欢运...\n",
      "\n",
      "=== 会话 'alice' 的记忆 ===\n",
      "用户 1: 我是Alice，我喜欢读书\n",
      "助手 2: 你好，Alice！很高兴认识你。你喜欢读书，那一定读过很多有趣的书吧？我特别想知道你最近在读什么书，或者有没有特别喜欢的类型？比如小说、科幻、历史还是其他类型的书呢？\n",
      "用户 3: 我刚才说喜欢什么？\n",
      "助手 4: 你刚才说：“我是Alice，我喜欢读书。” 所以你刚才说的是你喜欢读书。😊\n",
      "\n",
      "=== 会话 'bob' 的记忆 ===\n",
      "用户 1: 我是Bob，我喜欢运动\n",
      "助手 2: 你好，Bob！很高兴认识你。喜欢运动是个非常好的习惯，能让人保持健康和活力。你平时都喜欢哪些运动呢？是喜欢跑步、打球，还是健身、游泳之类的？我很好奇你的运动经历和爱好！😊\n",
      "用户 3: 我叫什么名字？\n",
      "助手 4: 你叫Bob呀！之前你已经告诉过我了，说\"我是Bob，我喜欢运动\"。所以，我当然知道你的名字是Bob啦！😊 你是不是在测试我的记忆力呢？不过没关系，我很乐意和你聊天，无论你问什么问题，我都会认真回答的！\n",
      "\n",
      "清除Alice的记忆\n",
      "已清除会话 alice 的记忆\n",
      "\n",
      "=== 会话 'alice' 的记忆 ===\n",
      "无记录\n"
     ]
    }
   ],
   "source": [
    "# 修复后的简化长期记忆管理\n",
    "\n",
    "from typing import Annotated, Sequence\n",
    "from typing_extensions import TypedDict\n",
    "from langgraph.graph import StateGraph, START, END, add_messages\n",
    "from langchain_core.messages import HumanMessage, AIMessage\n",
    "import json\n",
    "import redis\n",
    "from langchain_community.chat_models import ChatTongyi\n",
    "\n",
    "\n",
    "# 状态定义\n",
    "class GraphState(TypedDict):\n",
    "    messages: Annotated[Sequence, add_messages]\n",
    "\n",
    "\n",
    "# 大模型\n",
    "llm = ChatTongyi(\n",
    "    model_name=\"qwen-turbo\",\n",
    "    temperature=0.7,\n",
    "    streaming=True\n",
    ")\n",
    "\n",
    "\n",
    "# 长期记忆管理类\n",
    "class LongTermMemory:\n",
    "    def __init__(self, session_id: str, url: str = \"redis://localhost:6379/0\"):\n",
    "        self.session_id = session_id\n",
    "        self.key = f\"memory:{session_id}\"\n",
    "        try:\n",
    "            self.redis_client = redis.from_url(url)\n",
    "            self.redis_client.ping()\n",
    "        except Exception as e:\n",
    "            print(f\"Redis连接失败: {e}\")\n",
    "            self.redis_client = None\n",
    "\n",
    "    def save_memory(self, memory_type: str, content: str):\n",
    "        \"\"\"保存记忆\"\"\"\n",
    "        if not self.redis_client:\n",
    "            return\n",
    "        try:\n",
    "            memory_data = {\n",
    "                \"type\": memory_type,\n",
    "                \"content\": content\n",
    "            }\n",
    "            self.redis_client.lpush(self.key, json.dumps(memory_data))\n",
    "            print(f\"记忆已保存: {content[:30]}...\")\n",
    "        except Exception as e:\n",
    "            print(f\"保存记忆失败: {e}\")\n",
    "\n",
    "    def get_memory(self, limit: int = 10) -> list:\n",
    "        \"\"\"通过ID获取记忆\"\"\"\n",
    "        if not self.redis_client:\n",
    "            return []\n",
    "        try:\n",
    "            items = self.redis_client.lrange(self.key, 0, limit-1)\n",
    "            memories = []\n",
    "            for item in reversed(items):\n",
    "                data = json.loads(item)\n",
    "                if data[\"type\"] == \"human\":\n",
    "                    memories.append(HumanMessage(content=data[\"content\"]))\n",
    "                elif data[\"type\"] == \"ai\":\n",
    "                    memories.append(AIMessage(content=data[\"content\"]))\n",
    "            return memories\n",
    "        except Exception as e:\n",
    "            print(f\"获取记忆失败: {e}\")\n",
    "            return []\n",
    "\n",
    "    def clear_memory(self):\n",
    "        \"\"\"清除记忆\"\"\"\n",
    "        if not self.redis_client:\n",
    "            return\n",
    "        try:\n",
    "            self.redis_client.delete(self.key)\n",
    "            print(f\"已清除会话 {self.session_id} 的记忆\")\n",
    "        except Exception as e:\n",
    "            print(f\"清除记忆失败: {e}\")\n",
    "\n",
    "\n",
    "# 全局变量存储当前thread_id\n",
    "current_thread_id = None\n",
    "\n",
    "def chat_node(state: GraphState):\n",
    "    global current_thread_id\n",
    "    thread_id = current_thread_id or \"default\"\n",
    "    \n",
    "    # 初始化长期记忆\n",
    "    memory = LongTermMemory(session_id=thread_id)\n",
    "    \n",
    "    # 获取当前用户消息\n",
    "    user_message = state[\"messages\"][-1]\n",
    "    \n",
    "    if isinstance(user_message, HumanMessage):\n",
    "        # 保存用户消息到长期记忆\n",
    "        memory.save_memory(\"human\", user_message.content)\n",
    "    \n",
    "    # 获取历史记忆\n",
    "    historical_messages = memory.get_memory(limit=10)\n",
    "    \n",
    "    # 构建对话上下文\n",
    "    if historical_messages:\n",
    "        all_messages = historical_messages + [user_message]\n",
    "    else:\n",
    "        all_messages = [user_message]\n",
    "    \n",
    "    print(f\"使用 {len(all_messages)} 条消息作为上下文\")\n",
    "    \n",
    "    # 调用模型生成回复\n",
    "    response = llm.invoke(all_messages)\n",
    "    print(f\"AI回复: {response.content}\")\n",
    "    \n",
    "    # 保存AI回复到长期记忆\n",
    "    memory.save_memory(\"ai\", response.content)\n",
    "    \n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "# 构建图\n",
    "builder = StateGraph(GraphState)\n",
    "builder.add_node(\"chat\", chat_node)\n",
    "builder.add_edge(START, \"chat\")\n",
    "builder.add_edge(\"chat\", END)\n",
    "\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "app = builder.compile(checkpointer=MemorySaver())\n",
    "\n",
    "\n",
    "# 辅助函数\n",
    "def show_memory(thread_id: str):\n",
    "    \"\"\"显示指定ID的记忆\"\"\"\n",
    "    memory = LongTermMemory(session_id=thread_id)\n",
    "    messages = memory.get_memory(limit=20)\n",
    "    \n",
    "    print(f\"\\n=== 会话 '{thread_id}' 的记忆 ===\")\n",
    "    if not messages:\n",
    "        print(\"无记录\")\n",
    "    else:\n",
    "        for i, msg in enumerate(messages, 1):\n",
    "            role = \"用户\" if isinstance(msg, HumanMessage) else \"助手\"\n",
    "            print(f\"{role} {i}: {msg.content}\")\n",
    "\n",
    "\n",
    "def delete_memory(thread_id: str):\n",
    "    \"\"\"删除指定ID的记忆\"\"\"\n",
    "    memory = LongTermMemory(session_id=thread_id)\n",
    "    memory.clear_memory()\n",
    "\n",
    "\n",
    "# 测试长期记忆\n",
    "def test_long_term_memory():\n",
    "    global current_thread_id\n",
    "    \n",
    "    print(\"开始测试长期记忆系统\")\n",
    "    \n",
    "    # Alice 的对话\n",
    "    print(\"\\nAlice 第一次对话:\")\n",
    "    current_thread_id = \"alice\"\n",
    "    app.invoke(\n",
    "        {\"messages\": [HumanMessage(content=\"我是Alice，我喜欢读书\")]},\n",
    "        config={\"configurable\": {\"thread_id\": \"alice\"}}\n",
    "    )\n",
    "    \n",
    "    print(\"Alice 第二次对话:\")\n",
    "    current_thread_id = \"alice\"\n",
    "    app.invoke(\n",
    "        {\"messages\": [HumanMessage(content=\"我刚才说喜欢什么？\")]},\n",
    "        config={\"configurable\": {\"thread_id\": \"alice\"}}\n",
    "    )\n",
    "    \n",
    "    # Bob 的对话\n",
    "    print(\"\\nBob 第一次对话:\")\n",
    "    current_thread_id = \"bob\"\n",
    "    app.invoke(\n",
    "        {\"messages\": [HumanMessage(content=\"我是Bob，我喜欢运动\")]},\n",
    "        config={\"configurable\": {\"thread_id\": \"bob\"}}\n",
    "    )\n",
    "    \n",
    "    print(\"Bob 第二次对话:\")\n",
    "    current_thread_id = \"bob\"\n",
    "    app.invoke(\n",
    "        {\"messages\": [HumanMessage(content=\"我叫什么名字？\")]},\n",
    "        config={\"configurable\": {\"thread_id\": \"bob\"}}\n",
    "    )\n",
    "    \n",
    "    # 显示记忆\n",
    "    show_memory(\"alice\")\n",
    "    show_memory(\"bob\")\n",
    "    \n",
    "    # 清除Alice的记忆\n",
    "    print(\"\\n清除Alice的记忆\")\n",
    "    delete_memory(\"alice\")\n",
    "    show_memory(\"alice\")\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_long_term_memory()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MLOps",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
